﻿using System;
using System.Net;
using System.Net.Sockets;
using System.Threading.Tasks;

namespace TcpSocketTaskDemo
{
    /// <summary>
    /// TaskFactory.FromAsync is not suitable for socket, because its signature is not compatible
    /// with the "BeginXXX" method of the socket, which needs extra arguments
    /// if you want to use "FromAsync" on socket, you have to create NetworkStream each time, which 
    /// is rather troublesome
    /// so I make this class here, using TaskCompletionSource to return Task to represent the ongoing
    /// asynchronous receiving operation
    /// and the returned Task will not be Task<int> which returned by "EndReceive", which can only be useful
    /// when used together with other information, such as the receiving buffer
    /// but Task<byte[]> which contains the real data received by the socket
    /// </summary>
    sealed class AwaitSocket : IDisposable
    {
        // *************************************************************** //
        #region [ member fields ]

        private readonly Socket m_socket;
        private readonly byte[] m_recvBuffer;

        #endregion

        // *************************************************************** //
        #region [ constructor ]

        public AwaitSocket(Socket socket, int bufSize)
        {
            m_socket = socket;
            m_recvBuffer = new byte[bufSize];
        }

        #endregion

        // *************************************************************** //
        #region [ destructor ]

        private bool m_disposed;

        ~AwaitSocket()
        {
            Dispose(false);
        }

        public void Dispose()
        {
            Dispose(true);
            GC.SuppressFinalize(this);
        }

        private void Dispose(bool disposing)
        {
            if (m_disposed)
                return;

            try
            {
                if (disposing)
                {
                    m_socket.Shutdown(SocketShutdown.Both);
                    m_socket.Close();
                }
            }
            finally
            {
                m_disposed = true;
            }
        }

        #endregion

        // *************************************************************** //
        #region [ public APIs ]

        public Task ContinueConnect(IPEndPoint svrEndpoint)
        {
            TaskCompletionSource<object> tcs = new TaskCompletionSource<object>();

            TryConnect(tcs, svrEndpoint);

            return tcs.Task;
        }

        /// <summary>
        /// awaitable method which launch an asynchronous receiving operation
        /// and returns a token representing the completion
        /// </summary>
        /// <returns>
        /// token which represents the completion of the asynchronous receiving
        /// Tuple<byte[], int>.item1: original receiving buffer
        /// Tuple<byte[], int>.item2: valid length within the buffer that represents the meaningful data
        /// </returns>
        public Task<Tuple<byte[], int>> Receive()
        {
            TaskCompletionSource<Tuple<byte[], int>> tcs = new TaskCompletionSource<Tuple<byte[], int>>();

            AsyncCallback recvedCallback = (asyncResult) =>
            {
                try
                {
                    int recvLength = m_socket.EndReceive(asyncResult);
                    tcs.SetResult(Tuple.Create(m_recvBuffer, recvLength));
                }
                catch (Exception ex)
                {
                    tcs.SetException(ex);
                }
            };
            m_socket.BeginReceive(m_recvBuffer, 0, m_recvBuffer.Length, SocketFlags.None, recvedCallback, null);

            return tcs.Task;
        }

        public Task<int> Send(byte[] sendBuffer)
        {
            TaskCompletionSource<int> tcs = new TaskCompletionSource<int>();

            AsyncCallback sentCallback = (asyncResult) =>
            {
                try
                {
                    int sentLength = m_socket.EndReceive(asyncResult);
                    if (sentLength != sendBuffer.Length)
                        throw new NotImplementedException("always assume that data can be sent out in one shot");
                    tcs.SetResult(sentLength);
                }
                catch (Exception ex)
                {
                    tcs.SetException(ex);
                }
            };
            m_socket.BeginSend(sendBuffer, 0, sendBuffer.Length, SocketFlags.None, sentCallback, null);

            return tcs.Task;
        }

        #endregion

        // *************************************************************** //
        #region [ private helpers ]

        private void TryConnect(TaskCompletionSource<object> tcs, IPEndPoint svrEndpoint)
        {
            AsyncCallback connectCallback = (asyncResult) =>
            {
                try
                {
                    m_socket.EndConnect(asyncResult);
                    tcs.SetResult(null);
                }
                catch (SocketException sockError)
                {
                    if (sockError.SocketErrorCode == SocketError.ConnectionRefused)
                    {
                        Console.WriteLine("connection refused, try again,......");
                        TryConnect(tcs, svrEndpoint);
                    }
                    else
                        tcs.SetException(sockError);
                }
                catch (Exception ex)
                {
                    tcs.SetException(ex);
                }
            };
            m_socket.BeginConnect(svrEndpoint, connectCallback, null);
        }

        #endregion
    }
}
